import argparse
from collections import defaultdict

def read_bed(input_file):
    sites=defaultdict(list)
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            sites[(words[0],words[5])].append(line)
    return sites

def filter(sites,outfilefh,non_outfilefh,kmer=21):
    kmer=kmer//2
    for key in sites.keys():
        for i in range(len(sites[key])):
            iline=sites[key][i]
            iwords=iline.strip().split("\t")
            ipos=iwords[1]
            singleton_flag=1
            for j in range(max(0,i-kmer),min(i+kmer,len(sites[key]))):
                if i==j:
                    continue
                jline=sites[key][j]
                jwords=jline.strip().split("\t")
                jpos=jwords[1]
                #if int(jpos)==int(ipos):
                #    continue
                if int(jpos)<int(ipos)+kmer and int(jpos)>int(ipos)-kmer:
                    singleton_flag=0
                    break
                #if int(jpos)>int(ipos)+10:
                #    break
            if singleton_flag==0:
                non_outfilefh.write(iline)
            else:
                outfilefh.write(iline)
        

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/004_0111_001_R1_val_1_bismark_bt2_pe.deduplicated.bedGraph.gz.bismark.zero.cov.bed",
        help="input address.",
    )
    parser.add_argument(
        "--singleton",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/singleton.bisulfite",
        help="result address.",
    )
    parser.add_argument(
        "--nonsingleton",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite",
        help="result address.",
    )
    parser.add_argument(
        "--kmer",
        type=int,
        default=21,
        help="kmer.",
    )
    return parser.parse_args()

if __name__ == '__main__':
    args=parse_args()
    sites=read_bed(args.input)
    singletonfh = open(args.singleton+'.kmer.'+str(args.kmer)+'.bed', "w", buffering=512)
    nonsingletonfh = open(args.nonsingleton+'.kmer.'+str(args.kmer)+'.bed', "w", buffering=512)
    filter(sites,singletonfh,nonsingletonfh,args.kmer)